/*
*
给你一个下标从 0 开始的正整数数组 nums 和一个正整数 k 。

如果满足下述条件，则数对 (num1, num2) 是 优质数对 ：

num1 和 num2 都 在数组 nums 中存在。
num1 OR num2 和 num1 AND num2 的二进制表示中值为 1 的位数之和大于等于 k ，其中 OR 是按位 或 操作，而 AND 是按位 与 操作。
返回 不同 优质数对的数目。

如果 a != c 或者 b != d ，则认为 (a, b) 和 (c, d) 是不同的两个数对。例如，(1, 2) 和 (2, 1) 不同。

注意：如果 num1 在数组中至少出现 一次 ，则满足 num1 == num2 的数对 (num1, num2) 也可以是优质数对。

示例 1：

输入：nums = [1,2,3,1], k = 3
输出：5
解释：有如下几个优质数对：
- (3, 3)：(3 AND 3) 和 (3 OR 3) 的二进制表示都等于 (11) 。值为 1 的位数和等于 2 + 2 = 4 ，大于等于 k = 3 。
- (2, 3) 和 (3, 2)： (2 AND 3) 的二进制表示等于 (10) ，(2 OR 3) 的二进制表示等于 (11) 。值为 1 的位数和等于 1 + 2 = 3 。
- (1, 3) 和 (3, 1)： (1 AND 3) 的二进制表示等于 (01) ，(1 OR 3) 的二进制表示等于 (11) 。值为 1 的位数和等于 1 + 2 = 3 。
所以优质数对的数目是 5 。
示例 2：

输入：nums = [5,1,1], k = 10
输出：0
解释：该数组中不存在优质数对。

提示：

1 <= nums.length <= 105
1 <= nums[i] <= 109
1 <= k <= 60

  - @author ala
  - @date 2024-10-10 11:53
*/
package main

import (
	"fmt"
	"math/bits"
)

func main() {
	nums := []int{1, 2, 3, 1, 536870911}
	k := 3

	fmt.Println(countExcellentPairs(nums, k))
}

func countExcellentPairs(nums []int, k int) int64 {
	return V1(nums, k)
}

/**
 *	1）bc(x|y) + bc(x&y) = bc(x) + bc(y)
 *		其中：bc为BitCount
 *		把nums数组的每个元素去重，并且变为其32位二进制中1个个数，记为：数组bcs
 *		问题变为：(bcs中两数和 >=k 的组合数 * 2) + (单独一个数*2 >= k 的个数)
 *	2）第二部分在去重的同时就能统计出来
 *	3）第一部分用树状数组统计
 *		往左加，往右和
 */
func V1(nums []int, k int) int64 {
	//	去重，转换为BitCount
	ns, set := make([]int, 0, len(nums)), make(map[int]struct{}, len(nums))
	res := int64(0)
	for _, n := range nums {
		if _, ex := set[n]; !ex {
			oc := bits.OnesCount(uint(n))
			ns = append(ns, oc)

			if oc<<1 >= k {
				res += 1
			}
		}
		set[n] = struct{}{}
	}

	//	树状数组记录 和>= k 的数量
	bt := buildBitTree(32)
	for _, n := range ns {
		res += bt.sum(max(k-n, 1)) << 1
		bt.incr(n)
	}
	return res
}

type BitTree struct {
	tree []int64
	N    int
}

func buildBitTree(n int) *BitTree {
	bt := &BitTree{}
	bt.N = n
	bt.tree = make([]int64, n+1)
	return bt
}
func (bt *BitTree) incr(i int) {
	for j := i; j > 0; j -= j & -j {
		bt.tree[j] += 1
	}
}
func (bt *BitTree) sum(i int) int64 {
	sum := int64(0)
	for j := i; j <= bt.N; j += j & -j {
		sum += bt.tree[j]
	}
	return sum
}
